import asyncio
import logging
import time
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional, Dict, Any, List

from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncSession, 
    async_sessionmaker, 
    create_async_engine,
    AsyncEngine,
    AsyncConnection
)
from sqlalchemy.exc import SQLAlchemyError, DBAPIError, OperationalError
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from datetime import datetime

from app.config import settings
from app.utils.logging import get_logger

# 创建日志对象
logger = get_logger(__name__)

class Base(DeclarativeBase):
    """所有模型的基类"""
    # 类型映射配置
    type_annotation_map = {
        datetime: mapped_column(default=datetime.utcnow, nullable=False),
    }

class DatabaseManager:
    """数据库连接管理器，提供连接池、心跳检测和重连机制"""
    
    def __init__(self):
        self.engine: Optional[AsyncEngine] = None
        self.async_session_maker: Optional[async_sessionmaker[AsyncSession]] = None
        self._heartbeat_task: Optional[asyncio.Task] = None
        self._is_connected: bool = False
        self._last_heartbeat: float = 0
        self._reconnect_attempts: int = 0
        
    async def init_db(self) -> None:
        """初始化数据库连接"""
        try:
            logger.info(f"Initializing database connection to {settings.db.postgres_host}:{settings.db.postgres_port}")
            
            # 创建异步引擎
            self.engine = create_async_engine(
                settings.db.database_url,
                echo=settings.debug,  # 在调试模式下打印SQL语句
                pool_size=settings.db.db_pool_size,
                max_overflow=settings.db.db_max_overflow,
                pool_timeout=settings.db.db_pool_timeout,
                pool_recycle=settings.db.db_pool_recycle,
                pool_pre_ping=True,  # 自动检测断开的连接
            )
            
            # 创建会话工厂
            self.async_session_maker = async_sessionmaker(
                bind=self.engine,
                expire_on_commit=False,
                autoflush=False,
            )
            
            # 测试连接
            await self._test_connection()
            
            # 创建数据库表
            await self.create_tables()
            
            # 如果启用了心跳检测，启动心跳任务
            if settings.db.db_heartbeat_enabled:
                self._start_heartbeat_task()
                
            logger.info("Database connection initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize database: {str(e)}")
            raise
    
    async def _test_connection(self) -> None:
        """测试数据库连接"""
        try:
            async with self.engine.connect() as conn:
                result = await conn.execute(text("SELECT 1"))
                await conn.commit()
                
                if result.scalar_one() == 1:
                    self._is_connected = True
                    self._last_heartbeat = time.time()
                    self._reconnect_attempts = 0
                    logger.info("Database connection test successful")
                else:
                    self._is_connected = False
                    logger.error("Database connection test failed: unexpected result")
        except Exception as e:
            self._is_connected = False
            logger.error(f"Database connection test failed: {str(e)}")
            raise
    
    def _start_heartbeat_task(self) -> None:
        """启动心跳检测任务"""
        if self._heartbeat_task is None or self._heartbeat_task.done():
            logger.info(f"Starting database heartbeat task (interval: {settings.db.db_heartbeat_interval}s)")
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
    
    async def _heartbeat_loop(self) -> None:
        """心跳检测循环"""
        while True:
            try:
                await asyncio.sleep(settings.db.db_heartbeat_interval)
                await self._check_connection_health()
            except asyncio.CancelledError:
                logger.info("Database heartbeat task cancelled")
                break
            except Exception as e:
                logger.error(f"Error in database heartbeat loop: {str(e)}")
                await asyncio.sleep(5)  # 出错后等待短暂再重试
    
    async def _check_connection_health(self) -> None:
        """检查连接健康状态"""
        try:
            async with self.engine.connect() as conn:
                await conn.execute(text("SELECT 1"))
                await conn.commit()
                
                self._is_connected = True
                self._last_heartbeat = time.time()
                self._reconnect_attempts = 0
                logger.debug("Database heartbeat check successful")
        except Exception as e:
            self._is_connected = False
            logger.error(f"Database heartbeat check failed: {str(e)}")
            
            # 如果启用了自动重连，尝试重连
            if settings.db.db_reconnect_enabled:
                await self._try_reconnect()
    
    async def _try_reconnect(self) -> None:
        """尝试重新连接数据库"""
        if self._reconnect_attempts >= settings.db.db_reconnect_max_retries:
            logger.error(f"Failed to reconnect to database after {self._reconnect_attempts} attempts")
            return
        
        self._reconnect_attempts += 1
        logger.info(f"Attempting to reconnect to database (attempt {self._reconnect_attempts}/{settings.db.db_reconnect_max_retries})")
        
        try:
            # 等待重连间隔
            await asyncio.sleep(settings.db.db_reconnect_retry_interval)
            
            # 重新初始化数据库连接
            await self.init_db()
            
            logger.info("Successfully reconnected to database")
        except Exception as e:
            logger.error(f"Reconnection attempt failed: {str(e)}")
            
            # 如果还有重试次数，安排下一次重试
            if self._reconnect_attempts < settings.db.db_reconnect_max_retries:
                logger.info(f"Scheduling next reconnection attempt in {settings.db.db_reconnect_retry_interval} seconds")
                asyncio.create_task(self._try_reconnect())
    
    async def create_tables(self) -> None:
        """创建数据库表"""
        try:
            # 导入所有模型类以确保它们被注册
            from app.models.models import School, Class, Course, User, SystemLog
            
            logger.info("Creating database tables if not exist")
            async with self.engine.begin() as conn:
                # 创建所有表
                await conn.run_sync(Base.metadata.create_all)
            
            logger.info("Database tables created successfully")
        except Exception as e:
            logger.error(f"Failed to create database tables: {str(e)}")
            raise
    
    async def close(self) -> None:
        """关闭数据库连接"""
        logger.info("Closing database connections")
        
        # 取消心跳任务
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass
        
        # 关闭引擎
        if self.engine:
            await self.engine.dispose()
            self.engine = None
            self._is_connected = False
            logger.info("Database connections closed")
    
    @asynccontextmanager
    async def session(self) -> AsyncGenerator[AsyncSession, None]:
        """创建数据库会话的上下文管理器"""
        if not self.async_session_maker:
            raise RuntimeError("Database not initialized. Call init_db() first.")
        
        session = self.async_session_maker()
        start_time = time.time()
        try:
            yield session
            await session.commit()
        except Exception as e:
            await session.rollback()
            logger.error(f"Database session error: {str(e)}")
            raise
        finally:
            await session.close()
            duration = time.time() - start_time
            logger.debug(f"Database session completed in {duration:.3f}s")

# 全局数据库管理器实例
db_manager = DatabaseManager()

# 依赖注入函数
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """依赖注入函数，用于FastAPI路由中获取数据库会话"""
    if not db_manager._is_connected:
        logger.warning("Database connection not established, attempting to initialize")
        await db_manager.init_db()
    
    async with db_manager.session() as session:
        yield session
